#pragma once

#include "../linear_layouts_core/linear_layout.cuh"
#include "../linear_layouts_core/layout_operations.cuh"

namespace linear_layouts {

/**
 * Swizzled Layout Implementation
 * 
 * Swizzled layouts reorder memory addresses to reduce bank conflicts
 * and improve memory throughput. The Linear Layouts paper defines
 * MMA swizzling and provides algorithms for optimal swizzling.
 * 
 * Swizzling Pattern: ((i/per_phase) % max_phase) ⊕ (j/vec)) * vec ⊕ (j % vec)
 * where ⊕ is XOR, i,j are coordinates, and vec/per_phase/max_phase are parameters
 */

enum class SwizzleType {
    MMA_SWIZZLE,     // MMA-specific swizzling for Tensor Cores
    OPTIMAL_SWIZZLE, // Computed optimal swizzling 
    CUSTOM_SWIZZLE,  // User-defined swizzling pattern
    IDENTITY         // No swizzling (unswizzled layout)
};

template<SwizzleType TYPE = SwizzleType::MMA_SWIZZLE,
         int OUTPUT_DIMS = 16, int REG_BITS = 8, int THREAD_BITS = 10, int WARP_BITS = 5>
class SwizzledLayout {
public:
    static constexpr int INPUT_DIMS = REG_BITS + THREAD_BITS + WARP_BITS;
    using LayoutType = LinearLayout<INPUT_DIMS, OUTPUT_DIMS, REG_BITS, THREAD_BITS, WARP_BITS>;

private:
    LayoutType layout_;
    
    struct SwizzleConfig {
        int vec;          // Vectorization factor (power of 2)
        int per_phase;    // Elements per phase (power of 2) 
        int max_phase;    // Maximum phases (power of 2)
        int bank_bits;    // Number of bank address bits
        int element_bitwidth; // Bits per element
        int tensor_width; // Width of tensor being swizzled
        int tensor_height; // Height of tensor being swizzled
    };
    
    SwizzleConfig config_;

public:
    /**
     * Constructor for swizzled layout
     * 
     * @param tensor_width: Width of the tensor to be swizzled
     * @param tensor_height: Height of the tensor to be swizzled  
     * @param element_bitwidth: Bits per element (8, 16, 32, etc.)
     * @param vec: Vectorization factor (default auto-computed)
     * @param per_phase: Elements per swizzle phase (default auto-computed)
     * @param max_phase: Maximum swizzle phases (default auto-computed)
     */
    __host__ __device__ SwizzledLayout(
        int tensor_width = 16, int tensor_height = 16,
        int element_bitwidth = 16,
        int vec = 0, int per_phase = 0, int max_phase = 0) {
        
        config_.tensor_width = tensor_width;
        config_.tensor_height = tensor_height;
        config_.element_bitwidth = element_bitwidth;
        config_.bank_bits = 5; // 32 banks typical for GPU shared memory
        
        // Auto-compute swizzle parameters if not provided
        if (vec == 0 || per_phase == 0 || max_phase == 0) {
            compute_optimal_parameters();
        } else {
            config_.vec = vec;
            config_.per_phase = per_phase;
            config_.max_phase = max_phase;
        }
        
        // Validate parameters
        assert((config_.vec & (config_.vec - 1)) == 0); // Power of 2
        assert((config_.per_phase & (config_.per_phase - 1)) == 0);
        assert((config_.max_phase & (config_.max_phase - 1)) == 0);
        
        // Build the swizzled layout matrix
        build_swizzle_matrix();
    }

private:
    __host__ __device__ void compute_optimal_parameters() {
        // Compute optimal swizzling parameters based on tensor size and element width
        // This implements the optimal swizzling algorithm from the paper
        
        // Vectorization factor: maximize based on memory access patterns
        int bytes_per_element = (config_.element_bitwidth + 7) / 8;
        int max_vector_bytes = 16; // 128-bit vectorized loads on modern GPUs
        config_.vec = std::min(max_vector_bytes / bytes_per_element, config_.tensor_width);
        config_.vec = 1 << (__builtin_clz(config_.vec) ^ 31); // Round down to power of 2
        
        // Bank conflict parameters for shared memory
        int bank_width_bytes = 4; // 32-bit banks
        int banks_per_vector = (config_.vec * bytes_per_element + bank_width_bytes - 1) / bank_width_bytes;
        
        // Per-phase parameter: balance between bank conflict reduction and implementation complexity
        config_.per_phase = std::max(8, banks_per_vector * 2);
        config_.per_phase = 1 << (__builtin_clz(config_.per_phase) ^ 31);
        
        // Max phase: minimize bank conflicts across different access patterns
        config_.max_phase = std::min(8, config_.tensor_height / config_.per_phase);
        config_.max_phase = std::max(1, config_.max_phase);
        config_.max_phase = 1 << (__builtin_clz(config_.max_phase) ^ 31);
        
        // Adjust for specific swizzle types
        if constexpr (TYPE == SwizzleType::MMA_SWIZZLE) {
            // MMA swizzling uses specific parameters optimized for Tensor Cores
            config_.vec = 4;        // 4 elements per vector for FP16
            config_.per_phase = 8;  // 8 elements per phase
            config_.max_phase = 4;  // 4 phases maximum
        }
    }

    __host__ __device__ void build_swizzle_matrix() {
        layout_ = LayoutType();
        
        if constexpr (TYPE == SwizzleType::MMA_SWIZZLE) {
            build_mma_swizzle();
        } else if constexpr (TYPE == SwizzleType::OPTIMAL_SWIZZLE) {
            build_optimal_swizzle();
        } else if constexpr (TYPE == SwizzleType::IDENTITY) {
            build_identity_swizzle();
        }
    }

    __host__ __device__ void build_mma_swizzle() {
        // Implement MMA swizzling as defined in the paper:
        // offset = ((i/per_phase) % max_phase ⊕ j/vec) * vec ⊕ (j % vec)
        
        int vec_bits = __builtin_ctz(config_.vec);
        int phase_bits = __builtin_ctz(config_.max_phase);
        int per_phase_bits = __builtin_ctz(config_.per_phase);
        
        // Build the linear transformation matrix for the swizzle pattern
        // This implements the XOR operations as linear combinations over F₂
        
        int output_bit = 0;
        
        // Lower bits: j % vec (direct mapping)
        for (int bit = 0; bit < vec_bits; bit++) {
            layout_.template set_dimension_bits<Dimension::REG>(output_bit++, bit);
        }
        
        // Middle bits: (i/per_phase) % max_phase ⊕ j/vec
        for (int bit = 0; bit < phase_bits && output_bit < OUTPUT_DIMS; bit++) {
            // j/vec component
            if (vec_bits + bit < REG_BITS) {
                layout_.template set_dimension_bits<Dimension::REG>(output_bit, vec_bits + bit);
            } else if (vec_bits + bit < REG_BITS + THREAD_BITS) {
                layout_.template set_dimension_bits<Dimension::THREAD>(output_bit, vec_bits + bit - REG_BITS);
            }
            
            // i/per_phase component (XOR with j/vec)
            int i_bit_source = per_phase_bits + bit;
            if (i_bit_source < OUTPUT_DIMS/2) {
                // Map i bits to second half of output dimensions
                layout_.template set_dimension_bits<Dimension::REG>(OUTPUT_DIMS/2 + output_bit, i_bit_source);
                
                // XOR pattern: same output bit gets contributions from both i and j
                if (vec_bits + bit < REG_BITS) {
                    layout_.template set_dimension_bits<Dimension::REG>(OUTPUT_DIMS/2 + output_bit, vec_bits + bit);
                }
            }
            
            output_bit++;
        }
        
        // Remaining bits: direct mapping
        while (output_bit < OUTPUT_DIMS) {
            if (output_bit < REG_BITS) {
                layout_.template set_dimension_bits<Dimension::REG>(output_bit, output_bit);
            } else if (output_bit < REG_BITS + THREAD_BITS) {
                layout_.template set_dimension_bits<Dimension::THREAD>(output_bit, output_bit - REG_BITS);
            } else {
                layout_.template set_dimension_bits<Dimension::WARP>(output_bit, output_bit - REG_BITS - THREAD_BITS);
            }
            output_bit++;
        }
    }

    __host__ __device__ void build_optimal_swizzle() {
        // Implement the optimal swizzling algorithm from the paper
        // This maximizes vectorization while minimizing bank conflicts
        
        // Start with identity layout
        layout_ = LayoutType::identity();
        
        // Apply swizzling transformations based on computed parameters
        int vec_bits = __builtin_ctz(config_.vec);
        int bank_bits = config_.bank_bits;
        
        // Create swizzle pattern for bank conflict reduction
        for (int bank_bit = 0; bank_bit < bank_bits && bank_bit < OUTPUT_DIMS; bank_bit++) {
            for (int vec_bit = 0; vec_bit < vec_bits && vec_bit < OUTPUT_DIMS; vec_bit++) {
                // XOR pattern between bank and vector bits
                int output_row = vec_bits + bank_bit;
                if (output_row < OUTPUT_DIMS) {
                    layout_.template set_dimension_bits<Dimension::REG>(output_row, vec_bit);
                }
            }
        }
    }

    __host__ __device__ void build_identity_swizzle() {
        // No swizzling - identity mapping
        layout_ = LayoutType::identity();
    }

public:
    // Get the underlying linear layout
    __host__ __device__ const LayoutType& get_layout() const {
        return layout_;
    }
    
    __host__ __device__ LayoutType& get_layout() {
        return layout_;
    }

    // Apply swizzling to coordinates
    __host__ __device__ TensorCoord apply_swizzle(const TensorCoord& coord) const {
        HardwareCoord hw_coord(0, coord.x, coord.y);
        uint64_t swizzled = layout_.apply(hw_coord);
        
        // Decode swizzled coordinates
        int x_mask = (1 << __builtin_ctz(config_.tensor_width)) - 1;
        int y_mask = (1 << __builtin_ctz(config_.tensor_height)) - 1;
        
        int swizzled_x = static_cast<int>(swizzled & x_mask);
        int swizzled_y = static_cast<int>((swizzled >> __builtin_ctz(config_.tensor_width)) & y_mask);
        
        return TensorCoord(swizzled_x, swizzled_y);
    }

    // Compute memory offset with swizzling
    __host__ __device__ int swizzled_offset(int i, int j) const {
        if constexpr (TYPE == SwizzleType::MMA_SWIZZLE) {
            // Direct implementation of MMA swizzle formula
            int phase = (i / config_.per_phase) % config_.max_phase;
            int vec_group = j / config_.vec;
            int vec_offset = j % config_.vec;
            
            return ((phase ^ vec_group) * config_.vec) + vec_offset;
        } else {
            // Use the linear layout transformation
            auto swizzled_coord = apply_swizzle(TensorCoord(j, i));
            return swizzled_coord.y * config_.tensor_width + swizzled_coord.x;
        }
    }

    // Get swizzle configuration
    __host__ __device__ const SwizzleConfig& get_config() const {
        return config_;
    }

    // Compute bank conflicts for given access pattern
    __host__ __device__ int compute_bank_conflicts(const int* access_pattern, int num_accesses) const {
        int conflicts = 0;
        int bank_mask = (1 << config_.bank_bits) - 1;
        bool bank_used[32] = {false}; // Assume max 32 banks
        
        for (int i = 0; i < num_accesses; i++) {
            int offset = swizzled_offset(access_pattern[i * 2], access_pattern[i * 2 + 1]);
            int bank = offset & bank_mask;
            
            if (bank_used[bank]) {
                conflicts++;
            } else {
                bank_used[bank] = true;
            }
        }
        
        return conflicts;
    }

    // Check if swizzling improves memory throughput
    __host__ __device__ bool improves_throughput(const SwizzledLayout& unswizzled) const {
        // Compare vectorization factors and bank conflict potential
        int our_vector = config_.vec;
        int their_vector = unswizzled.get_config().vec;
        
        // Prefer higher vectorization
        if (our_vector > their_vector) return true;
        if (our_vector < their_vector) return false;
        
        // For equal vectorization, prefer better bank conflict reduction
        return config_.max_phase > unswizzled.get_config().max_phase;
    }

    // Create optimized swizzle for specific distributed layouts
    template<typename DistributedLayout>
    __host__ __device__ static SwizzledLayout create_for_layout(
        const DistributedLayout& read_layout,
        const DistributedLayout& write_layout) {
        
        // Analyze the distributed layouts to determine optimal swizzle parameters
        auto read_threads = read_layout.template extract_sublayout<Dimension::THREAD>();
        auto write_threads = write_layout.template extract_sublayout<Dimension::THREAD>();
        
        // Compute vectorization intersection
        int max_vec = 16; // Start with maximum possible vectorization
        
        // Reduce vectorization based on layout constraints
        if (!read_layout.supports_vectorized_load(max_vec)) max_vec /= 2;
        if (!write_layout.supports_vectorized_load(max_vec)) max_vec /= 2;
        
        return SwizzledLayout<SwizzleType::OPTIMAL_SWIZZLE>(16, 16, 16, max_vec, 8, 4);
    }

    // Convert between different swizzle patterns
    template<SwizzleType TARGET_TYPE>
    __host__ __device__ auto convert_to() const {
        SwizzledLayout<TARGET_TYPE, OUTPUT_DIMS, REG_BITS, THREAD_BITS, WARP_BITS> target(
            config_.tensor_width, config_.tensor_height, config_.element_bitwidth);
        return compose(target.get_layout().inverse(), layout_);
    }

    // Check compatibility with hardware primitives
    __host__ __device__ bool supports_ldmatrix() const {
        // ldmatrix requires specific swizzle patterns for optimal performance
        return config_.vec >= 4 && config_.per_phase >= 8;
    }

    __host__ __device__ bool supports_cp_async() const {
        // cp.async benefits from reduced bank conflicts
        return config_.max_phase >= 2;
    }

    // Estimate memory bandwidth utilization
    __host__ __device__ float estimated_bandwidth_utilization() const {
        // Simplified model: vectorization increases utilization, conflicts decrease it
        float vector_factor = static_cast<float>(config_.vec) / 16.0f; // Normalize to max
        float conflict_factor = 1.0f - (static_cast<float>(config_.max_phase - 1) / config_.max_phase);
        
        return vector_factor * conflict_factor;
    }

    // Create common swizzle configurations
    __host__ __device__ static SwizzledLayout create_mma_swizzle_fp16() {
        return SwizzledLayout<SwizzleType::MMA_SWIZZLE>(16, 16, 16, 4, 8, 4);
    }

    __host__ __device__ static SwizzledLayout create_mma_swizzle_fp8() {
        return SwizzledLayout<SwizzleType::MMA_SWIZZLE>(16, 16, 8, 8, 8, 4);
    }

    __host__ __device__ static SwizzledLayout create_optimal_for_shape(int width, int height, int bitwidth) {
        return SwizzledLayout<SwizzleType::OPTIMAL_SWIZZLE>(width, height, bitwidth);
    }

    // Debug information
    __host__ void print_info() const {
        printf("Swizzled Layout Configuration:\n");
        printf("  Type: ");
        if constexpr (TYPE == SwizzleType::MMA_SWIZZLE) printf("MMA_SWIZZLE\n");
        else if constexpr (TYPE == SwizzleType::OPTIMAL_SWIZZLE) printf("OPTIMAL_SWIZZLE\n");
        else if constexpr (TYPE == SwizzleType::IDENTITY) printf("IDENTITY\n");
        
        printf("  Tensor dimensions: %dx%d\n", config_.tensor_width, config_.tensor_height);
        printf("  Element bitwidth: %d\n", config_.element_bitwidth);
        printf("  Vectorization factor: %d\n", config_.vec);
        printf("  Elements per phase: %d\n", config_.per_phase);
        printf("  Maximum phases: %d\n", config_.max_phase);
        printf("  Bank bits: %d\n", config_.bank_bits);
        printf("  Supports ldmatrix: %s\n", supports_ldmatrix() ? "Yes" : "No");
        printf("  Supports cp.async: %s\n", supports_cp_async() ? "Yes" : "No");
        printf("  Estimated bandwidth utilization: %.2f\n", estimated_bandwidth_utilization());
        printf("  Matrix representation:\n");
        layout_.print();
    }

    // Test swizzle pattern with example coordinates
    __host__ void test_swizzle_pattern(int test_size = 8) const {
        printf("Swizzle pattern test (%dx%d):\n", test_size, test_size);
        printf("Original -> Swizzled\n");
        
        for (int i = 0; i < test_size; i++) {
            for (int j = 0; j < test_size; j++) {
                int original_offset = i * test_size + j;
                int swizzled_off = swizzled_offset(i, j);
                printf("(%d,%d):%d -> %d  ", i, j, original_offset, swizzled_off);
                
                if ((j + 1) % 4 == 0) printf("\n");
            }
        }
        printf("\n");
    }
};

// Common swizzled layout type aliases
using MMASwizzle16x16 = SwizzledLayout<SwizzleType::MMA_SWIZZLE, 16>;
using OptimalSwizzle16x16 = SwizzledLayout<SwizzleType::OPTIMAL_SWIZZLE, 16>;
using IdentityLayout16x16 = SwizzledLayout<SwizzleType::IDENTITY, 16>;

} // namespace linear_layouts